from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier

def main():
    iris = load_iris()  # 加载鸢尾花数据集
    x = iris.data
    y = iris.target
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.3, random_state=1)  # 数据集划分
    clf = DecisionTreeClassifier(criterion="gini")  # criterion="entropy" 特征=信息增益 gini
    clf.fit(x_train, y_train)
    train_score = clf.score(x_train, y_train)
    test_score = clf.score(x_test, y_test)
    print(f'train score: {train_score: .6f};\n test score: {test_score: .6f}')

    # 可视化处理
    # 1）使用 matplotlib.pyplot 进行绘制

    from sklearn.tree import plot_tree  # 导入决策树分类器及plot_tree函数绘制决策树
    import matplotlib.pyplot as plt
    iris_decision_tree1 = plt.figure(figsize=(9, 7), dpi=150)  # 设置图形大小
    plot_tree(clf, filled=True,  # 表示用不同颜色填充节点以表示不同的类别，
              feature_names=iris.feature_names,
              class_names=iris.target_names,
              rounded=True)  # 使节点的边框为圆形
    print(iris_decision_tree1)
    plt.show()  # 显示图片

    # 2）使用 graphviz 进行绘制

    # from sklearn.tree import export_graphviz  # 用于将训练好的决策树模型导出为 Graphviz 格式的图形描述文件的函数。
    # import graphviz
    # iris_decision_tree2 = export_graphviz(clf, out_file=None,  # 输出文件的路径，如果为None，则返回图形描述字符串。
    #                            feature_names=iris.feature_names,
    #                            class_names=iris.target_names,
    #                            filled=True, rounded=True,  # 表示用不同颜色填充节点以表示不同的类别，使节点的边框为圆形。
    #                            special_characters=True)  # 允许在节点标签中使用特殊字符。
    # tree = graphviz.Source(iris_decision_tree2, format='png')  # 将图形描述转换为图形对象,格式为png，默认为pdf
    # tree.render("iris_decision_tree2", view=True)  # render方法将图形保存为文件，view=True自动打开查看生成的图形文件。

if __name__ == '__main__':
    main()